#include "matplotlibcpp.h"
#include <cmath>
#include <iostream>
#include <random>
#include <ctime>

using namespace std;
namespace plt = matplotlibcpp;


int len_x = 1000, len_h = 500;
int N = 1000;
float de_coef = 0.0006;
float lambda = 1e-2;

//conv func
vector<float> conv2(vector<float> h, vector<float> x,int len_y){
    vector<float> y(len_y);
    for(int i=0; i<len_y; i++){
        for(int j=max(0,i+1-len_x); j<=min(i,len_h-1); j++){
            y[i] += h[j]*x[i-j];
        }
    }
    return y;
}

//noise creating func
vector<float> noise_cre(float mean, float std_dev, int len_y){
    default_random_engine generator;
    normal_distribution<float> dist(mean, std_dev);
    vector<float> noise(len_y);
    for(int i=0; i<len_y; i++)
        noise[i] = dist(generator);
    return noise;
}

vector<float> fliplr(vector<float> h){
    vector<float> H(len_h);
    for(int i=0; i<len_h; i++) {
        H[i] = h[len_h - i - 1];
    }
    return H;
}
int sign(float x){
    if(x>0)
        return 1;
    else if(x<0)
        return -1;
    else
        return 0;
}
//deconv with regularization
vector<float> deconv(vector<float> h,vector<float> y,int len_y){
    vector<float> x_k(len_x);
    for(int i=0; i<len_x; i++)
        x_k[i] = 0;
    vector<float> H = fliplr(h);
    for(int i=0; i<N; i++){
        vector<float> temp = conv2(h,x_k,len_y);
        for(int j=0; j < len_y; j++){
            temp[j] = temp[j] - y[j];
        }
        vector<float> delta_f = conv2(H,temp,len_y+len_h-1);
        vector<float> delta_d(len_x);
        for(int k=0; k < len_x; k++){
            delta_d[k] = delta_f[k + len_h -1];
        }

        vector<float> x_k1(len_x);
        for(int j=0; j<len_x; j++){
            x_k1[j] = x_k[j] - de_coef*delta_d[j];
        }

        for(int k=0; k<len_x; k++){
            if (abs(x_k1[k]) > lambda)
                x_k1[k] = sign(x_k1[k])*(abs(x_k1[k])-lambda);
            else
                x_k1[k] = 0;
        }

        x_k = x_k1;
    }

    return x_k;
}


int main() {
    clock_t  start,end;
    start = clock();

    vector<float> x(len_x);
    x[100] = x[300] = x[350] = x[500] = x[800] = 1;
    plt::plot(x);
    plt::figure();

    vector<float> h(len_h);
    for(int i=0; i<500; i++){
        h[i] = exp(-float(i)/50);
    }

    int len_y = len_x + len_h;
    vector<float> y = conv2(h, x, len_y);
    plt::plot(y);
    plt::figure();

    vector<float> noise = noise_cre(0.1, 0.1, len_y);
    vector<float> add_noise(len_y);
    for(int i=0; i<len_y; i++){
        add_noise[i] = y[i] + noise[i];
    }
    plt::plot(add_noise);
    plt::figure();


    vector<float> conv_x(len_x);
    conv_x = deconv(h,add_noise,len_y);
    plt::plot(conv_x);
    plt::show();

    end = clock();
    cout<<"total time = "<<(float)(end-start)*1000/CLOCKS_PER_SEC<<" ms"<<endl;

    return 0;

}

